package com.office.framework.websocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.office.common.utils.SecurityUtils;
import com.office.common.utils.redis.RedisUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 群聊WebSocket处理器
 *
 * @author 数字牧马人
 */
@Slf4j
@Component
public class ChatWebSocketHandler extends TextWebSocketHandler {

    @Autowired
    private ObjectMapper objectMapper;

    // 存储所有连接的会话
    private static final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    // 存储用户会话映射 - 修改为存储完整的用户信息
    private static final Map<String, ChatUserInfo> userSessionMap = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String sessionId = session.getId();
        sessions.put(sessionId, session);

        log.info("WebSocket连接建立: {}", sessionId);

        // 发送连接成功消息
        sendMessage(session, createSystemMessage("连接成功"));
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        try {
            String payload = message.getPayload();
            log.info("收到消息: {}", payload);

            ChatMessage chatMessage = objectMapper.readValue(payload, ChatMessage.class);

            switch (chatMessage.getType()) {
                case "join":
                    handleJoinMessage(session, chatMessage);
                    break;
                case "chat":
                    handleChatMessage(session, chatMessage);
                    break;
                case "leave":
                    handleLeaveMessage(session, chatMessage);
                    break;
                case "ai":
                    handleAIMessage(session, chatMessage);
                    break;
                default:
                    log.warn("未知消息类型: {}", chatMessage.getType());
            }
        } catch (Exception e) {
            log.error("处理WebSocket消息失败", e);
            sendMessage(session, createErrorMessage("消息处理失败"));
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String sessionId = session.getId();
        sessions.remove(sessionId);

        // 移除用户会话映射
        ChatUserInfo userInfo = userSessionMap.remove(sessionId);
        if (userInfo != null) {
            // 广播用户离开消息
            broadcastMessage(createLeaveMessage(userInfo.getUsername()));
        }

        log.info("WebSocket连接关闭: {}, 状态: {}", sessionId, status);
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("WebSocket传输错误", exception);
        String sessionId = session.getId();
        sessions.remove(sessionId);
        userSessionMap.remove(sessionId);
    }

    /**
     * 处理用户加入消息
     */
    private void handleJoinMessage(WebSocketSession session, ChatMessage message) {
        String sessionId = session.getId();
        String username = message.getUsername();
        String nickname = message.getNickname();
        String avatar = message.getAvatar();

        // 存储完整的用户信息
        ChatUserInfo userInfo = new ChatUserInfo(username, nickname, avatar, sessionId);
        userSessionMap.put(sessionId, userInfo);

        // 广播用户加入消息
        ChatMessage joinMessage = new ChatMessage();
        joinMessage.setType("join");
        joinMessage.setUsername(username);
        joinMessage.setNickname(nickname);
        joinMessage.setAvatar(avatar);
        joinMessage.setTimestamp(System.currentTimeMillis());

        broadcastMessage(joinMessage);

        log.info("用户加入: {}", username);
    }

    /**
     * 处理聊天消息
     */
    private void handleChatMessage(WebSocketSession session, ChatMessage message) {
        // 广播聊天消息
        broadcastMessage(message);

        log.info("用户 {} 发送消息: {}", message.getUsername(), message.getContent());
    }

    /**
     * 处理AI消息
     */
    private void handleAIMessage(WebSocketSession session, ChatMessage message) {
        // 广播AI消息
        broadcastMessage(message);

        log.info("AI助手回复: {}", message.getContent());
    }

    /**
     * 处理用户离开消息
     */
    private void handleLeaveMessage(WebSocketSession session, ChatMessage message) {
        String sessionId = session.getId();
        String username = message.getUsername();

        // 移除用户会话映射
        userSessionMap.remove(sessionId);

        // 广播用户离开消息
        broadcastMessage(message);

        log.info("用户离开: {}", username);
    }

    /**
     * 广播消息给所有连接的客户端
     */
    public void broadcastMessage(ChatMessage message) {
        String messageJson;
        try {
            messageJson = objectMapper.writeValueAsString(message);
        } catch (Exception e) {
            log.error("序列化消息失败", e);
            return;
        }

        TextMessage textMessage = new TextMessage(messageJson);

        sessions.values().forEach(session -> {
            try {
                if (session.isOpen()) {
                    session.sendMessage(textMessage);
                }
            } catch (IOException e) {
                log.error("发送消息失败", e);
            }
        });
    }

    /**
     * 发送消息给指定会话
     */
    private void sendMessage(WebSocketSession session, ChatMessage message) {
        try {
            String messageJson = objectMapper.writeValueAsString(message);
            TextMessage textMessage = new TextMessage(messageJson);
            session.sendMessage(textMessage);
        } catch (Exception e) {
            log.error("发送消息失败", e);
        }
    }

    /**
     * 创建系统消息
     */
    private ChatMessage createSystemMessage(String content) {
        ChatMessage message = new ChatMessage();
        message.setType("system");
        message.setContent(content);
        message.setTimestamp(System.currentTimeMillis());
        return message;
    }

    /**
     * 创建错误消息
     */
    private ChatMessage createErrorMessage(String content) {
        ChatMessage message = new ChatMessage();
        message.setType("error");
        message.setContent(content);
        message.setTimestamp(System.currentTimeMillis());
        return message;
    }

    /**
     * 创建离开消息
     */
    private ChatMessage createLeaveMessage(String username) {
        ChatMessage message = new ChatMessage();
        message.setType("leave");
        message.setUsername(username);
        message.setTimestamp(System.currentTimeMillis());
        return message;
    }

    /**
     * 获取当前在线用户数量
     */
    public int getOnlineUserCount() {
        // 使用用户名去重统计
        return (int) userSessionMap.values().stream()
                .map(ChatUserInfo::getUsername)
                .distinct()
                .count();
    }

    /**
     * 获取所有在线用户 - 返回完整的用户信息
     */
    public Map<String, ChatUserInfo> getOnlineUsers() {
        return new ConcurrentHashMap<>(userSessionMap);
    }
}
